"""
Author: Morphlng
Date: 2024-03-18 14:35:32
LastEditTime: 2024-12-16 23:34:31
LastEditors: Morphlng
Description: Scenario manager for all scenarios in one environment
FilePath: /DrivingGym/src/driving_gym/environment/scenario/scenario_manager.py
"""

from __future__ import annotations

import random
import os
from typing import TYPE_CHECKING

from driving_gym.environment.scenario.base_scenario import BaseScenario
from driving_gym.misc.util import dynamic_import

if TYPE_CHECKING:
    from driving_gym.simulation.adapter_interface import AdapterInterface


class ScenarioManager:
    def __init__(self, scenarios: "dict | list[dict]", adapter: "AdapterInterface"):
        """Create a scenario manager with a list of scenario configuration

        Args:
            scenarios (dict | list[dict]): dict(s) of scenario configuration
        """
        self.adapter = adapter

        if not isinstance(scenarios, list):
            scenarios = [scenarios]

        self.scenarios = scenarios
        self.idx = 0

        for i in range(len(self.scenarios)):
            validation, reason = self._validate_scenario(scenarios[i])
            if not validation:
                raise ValueError(
                    f"Scenario {i} is invalid. Missing required field `{reason}`"
                )

    def get_scenario(self, random_pick: bool = False) -> BaseScenario:
        """Load a scenario

        Args:
            random_pick (bool, optional): If True, pick a scenario randomly. Defaults to False.

        Returns:
            Scenario: A scenario instance
        """
        if random_pick:
            config = random.choice(self.scenarios)
        else:
            config = self.scenarios[self.idx]
            self.idx = (self.idx + 1) % len(self.scenarios)

        scenario_type = config.get("type", "base")
        root_dir = os.path.join(os.path.dirname(__file__))
        scenario = dynamic_import(scenario_type, root_dir, args=(config, self.adapter))
        return scenario

    def _validate_scenario(self, config: dict):
        """Validate a scenario configuration

        Args:
            config (dict): The scenario configuration

        Returns:
            bool, str: Whether the configuration is valid, and the missing field if not
        """
        if not "actor" in config:
            return False, "actor"

        if not "task" in config:
            return False, "task"

        if not "map" in config:
            return False, "map"

        return True, ""
